Rust-Linq Source

macro_rules! query {
    ([$p:pat] [$builder:expr] let $new_binding:pat = { $($new_value:tt)+ } $($rest:tt)*) => {
        query!([($p, $new_binding)] [$builder.map(|all @ $p| (all, { $($new_value)+ }))] $($rest)*)
    };

    ([$p:pat] [$builder:expr] where { $($filter:tt)+ } $($rest:tt)*) => {
        query!([$p] [$builder.filter(|$p| { $($filter)+ })] $($rest)*)
    };

    ([$p:pat] [$builder:expr] orderby { $($key:tt)+ } descending $($rest:tt)*) => {
        query!([$p] [Sorted::new($builder, |$p| { $($key)* })] $($rest)*)
    };

    ([$p:pat] [$builder:expr] orderby { $($key:tt)* } ascending $($rest:tt)*) => {
        query!([$p] [Sorted::new($builder, |$p| std::cmp::Reverse({ $($key)* }))] $($rest)*)
    };

    ([$p:pat] [$builder:expr] orderby { $($key:tt)* } $($rest:tt)*) => {
        query!([$p] [Sorted::new($builder, |$p| std::cmp::Reverse({ $($key)* }))] $($rest)*)
    };

    ([$p:pat] [$builder:expr] group { $($map:tt)+ } by { $($key:tt)+ } into { $group:pat } $($rest:tt)*) => {
        query!([$group] [Grouped::new($builder, |$p| { $($key)+ }, |$p| { $($map)+ })] $($rest)*)
    };

    ([$p:pat] [$builder:expr] from $new_binding:pat in { $($source:tt)+ } $($rest:tt)* ) => {
        query!([($p, $new_binding)] [Layered::new($builder, |$p| { $($source)+ })] $($rest)*)
    };

    ([$p:pat] [$builder:expr] select $($select:tt)+) => {
        $builder.map(|$p| { $($select)+ })
    };

    (from $p:pat in { $source:expr } $($rest:tt)*) => {
        query!([$p] [$source.into_iter()] $($rest)*)
    }
}

#[derive(Clone)]
pub enum Sorted<T, E, F> {
    Ready(Option<T>, F),
    Active(Vec<E>),
}

impl<I, E, F, K> Sorted<I, E, F>
where
    I: IntoIterator<Item = E>,
    F: FnMut(&E) -> K,
    K: Ord,
{
    fn new(items: I, key: F) -> Self {
        Self::Ready(Some(items), key)
    }
}

impl<I, E, F, K> Iterator for Sorted<I, E, F>
where
    I: IntoIterator<Item = E>,
    F: FnMut(&E) -> K,
    K: Ord,
{
    type Item = E;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::Ready(items, key) => {
                let mut items: Vec<_> = items.take()?.into_iter().collect();
                items.sort_unstable_by_key(key);
                *self = Self::Active(items);
                self.next()
            }
            Self::Active(items) => items.pop(),
        }
    }
}

#[derive(Clone)]
pub enum Grouped<T, FK, FV, K, V> {
    Ready(Option<T>, FK, FV),
    Active(Vec<(K, Vec<V>)>),
}

impl<I, E, FK, FV, K, V> Grouped<I, FK, FV, K, V>
where
    I: IntoIterator<Item = E>,
    FK: FnMut(&E) -> K,
    FV: FnMut(E) -> V,
    K: Eq + std::hash::Hash,
{
    fn new(items: I, key: FK, value: FV) -> Self {
        Self::Ready(Some(items), key, value)
    }
}

impl<I, E, FK, FV, K, V> Iterator for Grouped<I, FK, FV, K, V>
where
    I: IntoIterator<Item = E>,
    FK: FnMut(&E) -> K,
    FV: FnMut(E) -> V,
    K: Eq + std::hash::Hash,
{
    type Item = (K, Vec<V>);

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::Ready(items, key, value) => {
                let mut map = std::collections::HashMap::new();
                for item in items.take()? {
                    map.entry(key(&item)).or_insert(vec![]).push(value(item));
                }
                *self = Self::Active(map.into_iter().collect());
                self.next()
            }
            Self::Active(items) => items.pop(),
        }
    }
}

pub struct Layered<I1, E1, FI2, I2> {
    outer: I1,
    inner: FI2,
    current: Option<(E1, I2)>,
}

impl<I1, E1, FI2, I2, E2> Layered<I1, E1, FI2, I2>
where
    I1: Iterator<Item = E1>,
    E1: Clone,
    FI2: FnMut(&E1) -> I2,
    I2: Iterator<Item = E2>,
{
    fn new(outer: I1, inner: FI2) -> Self {
        Self {
            outer,
            inner,
            current: None,
        }
    }
}

impl<I1, E1, FI2, I2, E2> Iterator for Layered<I1, E1, FI2, I2>
where
    I1: Iterator<Item = E1>,
    E1: Clone,
    FI2: FnMut(&E1) -> I2,
    I2: Iterator<Item = E2>,
{
    type Item = (E1, E2);

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(ref mut current) = self.current {
            if let Some(item) = current.1.next() {
                Some((current.0.clone(), item))
            } else {
                self.current = None;
                self.next()
            }
        } else if let Some(next) = self.outer.next() {
            let next_inner = (self.inner)(&next);
            self.current = Some((next, next_inner));
            self.next()
        } else {
            None
        }
    }
}

fn main() {
    let y = [8, 0, 5, 9, 3, 6, 7, 1, 2, 4];

    #[allow(unused_variables)]
    let it = query! {
        from v in { y }
        where { *v > 5 }
        orderby { *v } ascending
        group { v } by { v % 2 == 0 } into { g }
        select g
    };

    for x in it {
        println!("{x:?}");
    }

    #[allow(unused_variables)]
    let it = query! {
        from a in { 0..10 }
        from b in { *a..10 }
        orderby { *b } ascending
        group { b } by { *a } into { g }
        orderby { g.0 } ascending
        select g.1
    };

    for x in it {
        println!("{x:?}");
    }
}